package io.github.qifan777.knowledge.ai.agent.common;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
import io.github.qifan777.knowledge.ai.agent.AbstractAgent;
import io.github.qifan777.knowledge.ai.agent.Agent;
import lombok.AllArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Description;
import org.springframework.http.*;
import org.springframework.web.client.RestClient;
import org.springframework.web.client.RestTemplate;

import java.util.*;
import java.util.function.Function;

/**
 * @author tjw
 * @date 2025-04-23
 * @Description
 */
@Configuration
@Agent
@Description("通用工具")
@AllArgsConstructor
public class CommonAssistant extends AbstractAgent implements Function<CommonAssistant.Request, String> {

    private final String SYSTEM = """
            通用工具
            """;

    private final ChatModel chatModel;

    private final String apiKey = "sk-798de318563045da9fba68af30d87925";


    @Override
    public String apply(Request request) {
        // 调用模型
        return ChatClient.create(chatModel)
                .prompt()
                // 系统提示
                .system(SYSTEM)
                //用户提示词
                .user(request.query)
                // 注册functioncalling，带有@bean的方法
                .functions("getWeatherByCity","generateImage")
                // call代表非流式问答，返回的结果可以是ChatResponse，也可以是Entity（转成java类型），也可以是字符串直接提取回答结果
                .call()
                .content();
    }

    /**
     * record是Java 14引入的一个新特性，用于创建不可变的数据载体类（也称为“记录类”）
     */
    public record Request(@JsonProperty(required = true) @JsonPropertyDescription("用户原始的提问") String query) {
    }


    public record WeatherRequest(String city) {
    }

    public record GenerateImageRequest(String prompt) {
    }

    /**
     * 图片生成 functioncalling
     * @return
     */
    @Bean
    @Description("获取生成的图片")
    public Function<CommonAssistant.GenerateImageRequest, List<String>> generateImage(){
        return request -> {
            String taskId = getTaskId(request);
            return getImageUrlByTaskId(taskId);
        };
    }

    /**
     *
     * 根据taskId查询图片生成结果
     * @param taskId
     * @return
     */
    public List<String> getImageUrlByTaskId(String taskId) {
        // 创建 RestTemplate 实例
        RestTemplate restTemplate = new RestTemplate();
        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "Bearer " + apiKey);
        // 根据taskId查询图片
        String url = "https://dashscope.aliyuncs.com/api/v1/tasks/" + taskId;
        // 创建 HttpEntity
        HttpEntity<String> requestEntity = new HttpEntity<>(headers);
        ResponseEntity<JSONObject> response;
        JSONObject jsonObject  ;
        JSONObject output = null;
        int i = 0;
        // 堵塞等待任务完成
        while (i < 30){
            response = restTemplate.exchange(url, HttpMethod.GET,
                    requestEntity ,JSONObject.class);
            jsonObject = response.getBody();
            output = jsonObject.getJSONObject("output");
            String taskStatus = output.getString("task_status");
            if (taskStatus.equals("SUCCEEDED")) {
                break;
            } else if (taskStatus.equals("FAILED")) {
                return null;
            }
            try {
                Thread.sleep(500);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            i++;
        }
        JSONArray results = output.getJSONArray("results");
        List<String> list = new ArrayList<>(results.size());
        for (Object result : results) {
            LinkedHashMap<String,String> item = (LinkedHashMap) result;
            list.add(item.get("url"));
        }
        return list;
    }

    /**
     * 调用通义万相-文生图模型创建生成图片任务
     * @param request
     * @return
     */
    public  String getTaskId(GenerateImageRequest request) {
        // 创建生成图片任务
        String url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis";
        RestTemplate restTemplate = new RestTemplate();
        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("X-DashScope-Async", "enable");
        headers.set("Authorization", "Bearer " + apiKey);
        // 创建请求体
        Map<String, Object> requestBody = new HashMap<>();
        Map<String, String> input = new HashMap<>();
        Map<String, Object> parameters = new HashMap<>();
        // 提示词
        input.put("prompt", request.prompt);
        // 图片大小
        parameters.put("size", "1024*1024");
        // 图片张数，取值1-4
        parameters.put("n", 1);
        // 图片模型
        requestBody.put("model", "wanx2.1-t2i-turbo");
        requestBody.put("input", input);
        requestBody.put("parameters", parameters);
        // 创建 HttpEntity
        HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);
        // 发起 POST 请求
        ResponseEntity<JSONObject> response = restTemplate.postForEntity(url, requestEntity, JSONObject.class);
        JSONObject jsonObject = response.getBody();
        JSONObject output = jsonObject.getJSONObject("output");
        return output.getString("task_id");
    }

    @Bean
    @Description("获取城市天气")
    public Function<CommonAssistant.WeatherRequest, JSONObject> getWeatherByCity(){
        return request -> {
            RestClient client = RestClient.create();
            return client.get()
                    .uri("https://api.vvhan.com/api/weather?city={0}", request.city)
                    .retrieve()
                    .body(JSONObject.class) ;
        };

    }
}
